# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 2023/4/24 22:53
# @Author : xyjiang1216
# @Email : jiangxuyan@nudt.edu.cn
# @File : frame_algo_utils.py
# @Version: OpenPlanner2.0
import math
import os
import re
import time

from z3 import And, Bool, Or, Solver, sat, unsat, unknown

from OpenPlanner.core.src.op_class import OPClass
from OpenPlanner.util.compute_hyper_period import compute_hyper_period


def instance_reinit_hook(op_obj: OPClass):
    stream_obj_set = op_obj.algo_input.stream_array
    link_obj_set = op_obj.algo_input.link_array
    stream_instance_obj_set = op_obj.algo_input.instance_array
    for stream_obj in stream_obj_set:
        period = stream_obj.period
        route = stream_obj.route
        hop_id = 0
        stream_id = stream_obj.stream_id
        size = stream_obj.size
        for link_id in route:
            macrotick = link_obj_set[link_id].macrotick
            speed = link_obj_set[link_id].speed
            trans_duration = math.ceil(size * 8 / speed)
            for instance_obj in stream_instance_obj_set[stream_id][hop_id]:
                instance_obj.period_scaled_to_macrotick = math.ceil(period / macrotick)
                instance_obj.trans_duration_scaled_to_macrotick = math.ceil(trans_duration /
                                                                            macrotick)
            hop_id += 1

    # resize instance obj set
    instance_obj_set_temp = []
    for instance_obj_set_per_stream in stream_instance_obj_set:
        instance_obj_set_per_link_temp = []
        for instance_obj_set_per_link in instance_obj_set_per_stream:
            instance_obj_set_per_link_temp.append(
                instance_obj_set_per_link[0]
            )
        instance_obj_set_temp.append(instance_obj_set_per_link_temp)
    # stream_instance_obj_set = instance_obj_set_temp
    op_obj.algo_input.instance_array = \
        instance_obj_set_temp


def construct_constraints(link_obj_set,
                          stream_obj_set,
                          stream_instance_obj_set,
                          sync_precision):
    constraint_formula_set = []

    # 1. frame constraint
    for stream_instance_obj_set_per_stream in stream_instance_obj_set:
        for stream_instance_obj in stream_instance_obj_set_per_stream:
            formula = And(stream_instance_obj.offset >= 0,
                          stream_instance_obj.offset <=
                          stream_instance_obj.period_scaled_to_macrotick -
                          stream_instance_obj.trans_duration_scaled_to_macrotick)
            constraint_formula_set.append(formula)

    # 2. link constraint
    for link in link_obj_set:
        stream_set = link.stream_array
        link_id = link.link_id
        for i in range(len(stream_set)):
            for j in range(len(stream_set)):
                if i != j:
                    i_stream_id = stream_set[i].stream_id
                    i_hop_id = stream_set[i].hop_id
                    j_stream_id = stream_set[j].stream_id
                    j_hop_id = stream_set[j].hop_id

                    ik_offset = stream_instance_obj_set[i_stream_id][i_hop_id].offset
                    ik_period_scaled_to_macrotick = stream_instance_obj_set[i_stream_id][
                        i_hop_id].period_scaled_to_macrotick
                    ik_trans_duration = stream_instance_obj_set[i_stream_id][
                        i_hop_id].trans_duration_scaled_to_macrotick
                    jl_offset = stream_instance_obj_set[j_stream_id][j_hop_id].offset
                    jl_period_scaled_to_macrotick = stream_instance_obj_set[j_stream_id][
                        j_hop_id].period_scaled_to_macrotick
                    jl_trans_duration = stream_instance_obj_set[j_stream_id][
                        j_hop_id].trans_duration_scaled_to_macrotick

                    i_period = stream_obj_set[i_stream_id].period
                    j_period = stream_obj_set[j_stream_id].period

                    i_j_hyper_period = compute_hyper_period(*[i_period, j_period])

                    for alpha in range(math.ceil(i_j_hyper_period / i_period)):
                        for beta in range(math.ceil(i_j_hyper_period / j_period)):
                            formula = Or(ik_offset + alpha * ik_period_scaled_to_macrotick >=
                                         jl_offset + beta * jl_period_scaled_to_macrotick +
                                         jl_trans_duration,
                                         jl_offset + beta * jl_period_scaled_to_macrotick >=
                                         ik_offset + alpha * ik_period_scaled_to_macrotick + ik_trans_duration)
                            constraint_formula_set.append(formula)

    # 3.flow transmission constraint
    for stream_instance_obj in stream_instance_obj_set:
        for i in range(len(stream_instance_obj) - 1):
            ax_link_id = stream_instance_obj[i].link_id
            xb_link_id = stream_instance_obj[i + 1].link_id
            xb_macrotick = link_obj_set[xb_link_id].macrotick
            ax_macrotick = link_obj_set[ax_link_id].macrotick
            # ax_propagation_delay = link_obj_set[ax_link_id].prop_delay
            xb_offset = stream_instance_obj[i + 1].offset
            ax_offset = stream_instance_obj[i].offset
            ax_trans_duration = stream_instance_obj[i].trans_duration_scaled_to_macrotick
            formula = (xb_offset * xb_macrotick -
                       sync_precision >=
                       (ax_offset + ax_trans_duration) * ax_macrotick)
            constraint_formula_set.append(formula)

    # 4. end-to-end constraint
    stream_id = 0
    for stream_instance_obj_set_per_stream in stream_instance_obj_set:
        latency_requirement = stream_obj_set[stream_id].latency_requirement
        src_offset = stream_instance_obj_set_per_stream[0].offset
        src_link_id = stream_instance_obj_set_per_stream[0].link_id
        src_macrotick = link_obj_set[src_link_id].macrotick
        dst_offset = stream_instance_obj_set_per_stream[-1].offset
        dst_link_id = stream_instance_obj_set_per_stream[-1].link_id
        dst_macrotick = link_obj_set[dst_link_id].macrotick
        trans_duration_scaled_to_macrotick = stream_instance_obj_set_per_stream[
            -1].trans_duration_scaled_to_macrotick

        formula = (src_macrotick * src_offset + latency_requirement >=
                   dst_macrotick * (dst_offset + trans_duration_scaled_to_macrotick))
        constraint_formula_set.append(formula)

        stream_id += 1

    # 6. Frame isolation constraint
    for link in link_obj_set:
        stream_set = link.stream_array
        link_id = link.link_id
        ab_macrotick = link_obj_set[link_id].macrotick
        # ab_prop_delay = link_obj_set[link_id].prop_delay
        for i in range(len(stream_set)):
            for j in range(len(stream_set)):
                if i != j:
                    i_stream_id = stream_set[i].stream_id
                    ik_ab_hop_id = stream_set[i].hop_id
                    j_stream_id = stream_set[j].stream_id
                    jl_ab_hop_id = stream_set[j].hop_id
                    if ik_ab_hop_id != 0 and jl_ab_hop_id != 0:

                        ik_xa_hop_id = stream_set[i].hop_id - 1
                        jl_ya_hop_id = stream_set[j].hop_id - 1

                        jl_ab_offset = stream_instance_obj_set[j_stream_id][jl_ab_hop_id].offset
                        jl_ya_offset = stream_instance_obj_set[j_stream_id][jl_ya_hop_id].offset

                        j_period = stream_obj_set[j_stream_id].period
                        i_period = stream_obj_set[i_stream_id].period

                        ik_xa_offset = stream_instance_obj_set[i_stream_id][ik_xa_hop_id].offset

                        ik_ab_offset = stream_instance_obj_set[i_stream_id][ik_ab_hop_id].offset

                        xa_link_id = stream_instance_obj_set[i_stream_id][ik_xa_hop_id].link_id
                        xa_macrotick = link_obj_set[xa_link_id].macrotick
                        # xa_prop_delay = link_obj_set[xa_link_id].prop_delay

                        ya_link_id = stream_instance_obj_set[j_stream_id][jl_ya_hop_id].link_id
                        ya_macrotick = link_obj_set[ya_link_id].macrotick
                        # ya_prop_delay = link_obj_set[ya_link_id].prop_delay

                        i_ab_prio = stream_instance_obj_set[i_stream_id][ik_ab_hop_id].prio
                        j_ab_prio = stream_instance_obj_set[j_stream_id][jl_ab_hop_id].prio

                        # formula = (Const(1, IntSort()) == 1)

                        formula = Bool('p')

                        i_j_hyper_period = compute_hyper_period(*[i_period, j_period])

                        for alpha in range(int(math.ceil(i_j_hyper_period / i_period))):
                            for beta in range(math.ceil(i_j_hyper_period / j_period)):
                                formula = And(formula,
                                              Or(
                                                  jl_ab_offset * ab_macrotick + beta * j_period
                                                  + sync_precision <=
                                                  ik_xa_offset * xa_macrotick +
                                                  alpha * i_period,
                                                  ik_ab_offset * ab_macrotick + alpha * i_period
                                                  + sync_precision <=
                                                  jl_ya_offset * ya_macrotick +
                                                  beta * j_period
                                              ))

                        formula = Or(formula,
                                     i_ab_prio != j_ab_prio)
                        # print(formula)
                        constraint_formula_set.append(formula)
                    else:
                        jl_ab_offset = stream_instance_obj_set[j_stream_id][jl_ab_hop_id].offset
                        j_period = stream_obj_set[j_stream_id].period
                        i_period = stream_obj_set[i_stream_id].period
                        ik_ab_offset = stream_instance_obj_set[i_stream_id][ik_ab_hop_id].offset
                        i_ab_prio = stream_instance_obj_set[i_stream_id][ik_ab_hop_id].prio
                        j_ab_prio = stream_instance_obj_set[j_stream_id][jl_ab_hop_id].prio
                        formula = Bool('p')

                        i_j_hyper_period = compute_hyper_period(*[i_period, j_period])

                        for alpha in range(int(math.ceil(i_j_hyper_period / i_period))):
                            for beta in range(math.ceil(i_j_hyper_period / j_period)):
                                formula = And(formula,
                                              Or(
                                                  jl_ab_offset * ab_macrotick + beta * j_period
                                                  + sync_precision <=
                                                  ik_ab_offset * ab_macrotick +
                                                  alpha * i_period,
                                                  ik_ab_offset * ab_macrotick + alpha * i_period
                                                  + sync_precision <=
                                                  jl_ab_offset * ab_macrotick +
                                                  beta * j_period
                                              ))
                        formula = Or(formula,
                                     i_ab_prio != j_ab_prio)

                        constraint_formula_set.append(formula)

    for stream_instance_set in stream_instance_obj_set:
        for stream_instance in stream_instance_set:
            link_id = stream_instance.link_id
            formula = (stream_instance.prio < link_obj_set[link_id].st_queues,
                       stream_instance.prio >= 0)
            constraint_formula_set.append(formula)
    return constraint_formula_set


def _parse_z3_model(model):
    solution = []
    for declare in model.decls():
        name = declare.name()
        value = model[declare]
        solution.append({'name': name, 'value': value})
    return solution


def add_and_solve_constraints(constraint_set,
                              timeout=-1):
    start = 0
    end = 0
    s = Solver()
    if timeout > 0:
        s.set(timeout=timeout)

    for constraint in constraint_set:
        s.add(constraint)

    declare_set = []
    unknown_reason = ''
    # 开始计时
    start = time.time_ns()
    # 判断是否有可行解
    sat_or_not = s.check()
    if sat_or_not == sat:
        model = s.model()
        end = time.time_ns()
        print("start time: %f" % start)
        print("end time: %f" % end)
        # 输出变量声明的集合
        declare_set = _parse_z3_model(model)
    elif sat_or_not == unsat:
        # 输出时间
        end = time.time_ns()
        # 输出一个空的declare_set
    elif sat_or_not == unknown:
        end = time.time_ns()
        # 输出一个空的declare_set
        # 输出unknown的原因
        unknown_reason = s.reason_unknown()
        pass

    time_used_in_second = (end - start) / 1000000000

    print('time_used:')
    print(time_used_in_second)

    print('sat_or_not:')
    print(str(sat_or_not))

    # 返回time_used_in_second、sat_or_not、
    # declare_set、(unknown reason)
    return {'time_used_in_second': time_used_in_second, 'sat_or_not': str(sat_or_not),
            'declare_set': declare_set, 'unknown_reason': unknown_reason}


def _classify_declare_set(declare_set, link_num):
    tmp_declare_set = [{'offset_set': [], 'prio_set': []} for link_id in range(link_num)]

    # frame_demo的变量分为两类
    # 1. offset，命名：O_stream_id^(link_id)
    # 2. prio，命名：P_stream_id^(link_id)
    for declare in declare_set:
        name = declare['name']
        value = declare['value']
        value = str(value)

        if re.match(r'O.', name):
            # 解析link_id
            link_id = int(name.split('(')[1].split(')')[0])
            declare['value'] = int(value)
            tmp_declare_set[link_id]['offset_set'].append(declare)
        elif re.match(r'P.', name):
            # 解析link_id
            link_id = int(name.split('(')[1].split(')')[0])
            declare['value'] = int(value)
            tmp_declare_set[link_id]['prio_set'].append(declare)

    declare_set = tmp_declare_set

    # 重排列
    # offset与prio应当一一对应
    # 依次取出每条链路的offset_set和prio_set
    for declare_set_per_link in declare_set:
        offset_set = declare_set_per_link['offset_set']
        prio_set = declare_set_per_link['prio_set']
        tmp_prio_set = [''] * len(offset_set)

        for prio in prio_set:
            name = prio['name']
            # 提取出stream_id
            prio_stream_id = name.split('_')[1].split('^')[0]
            idx = 0
            for offset in offset_set:
                offset_stream_id = offset['name'].split('_')[1].split('^')[0]
                if prio_stream_id == offset_stream_id:
                    break
                idx += 1
            # print(idx)
            tmp_prio_set[idx] = prio

        declare_set_per_link['prio_set'] = tmp_prio_set
    return declare_set


def _sort_declare_set(declare_set):
    # print(declare_set)
    for declare in declare_set:
        offset_set = declare['offset_set']
        prio_set = declare['prio_set']
        if offset_set and prio_set:
            # print(offset_set)
            # print(prio_set)
            zipped_set = zip(offset_set, prio_set)
            # print(zipped_set)
            sorted_zipped_set = sorted(zipped_set, key=lambda tmp: (tmp[0]['value'], tmp[1]['value']))
            sorted_set = zip(*sorted_zipped_set)
            # print(sorted_set)
            # for x in sorted_set:
            #     print(list(x))
            declare['offset_set'], declare['prio_set'] = [list(x) for x in sorted_set]
    # print(declare_set)
    return declare_set


def _format_offset_set_and_prio_set(f, declare, width):
    offset_set = declare['offset_set']
    prio_set = declare['prio_set']

    f.write('-' * (width - 1) + '|' + '\n')
    for offset, prio in zip(offset_set, prio_set):
        f.write('{:>15}{:^20}{}{:^15}{}\n'.format('offset:', offset['name'], '=', offset['value'], '|'))
        f.write('{:>15}{:^20}{}{:^15}{}\n'.format('prio_queues:', prio['name'], '=', prio['value'], '|'))
        f.write('-' * (width - 1) + '|' + '\n')


def _write_time_used_to_txt(f, time_used_in_second):
    divider = '+' * 80
    f.write('%s\n' % divider)
    f.write('time used:\n')
    f.write("%s s\n" % time_used_in_second)
    f.write("%s min\n" % (time_used_in_second / 60))
    return


def _expand_declare_set(z3_declare_set, link_obj_set, stream_obj_set):
    temp_z3_declare_set = [{'offset_set': [], 'prio_set': []} for declare in z3_declare_set]
    # print(temp_z3_declare_set)

    hyper_period = compute_hyper_period(
        *[stream_obj.period for stream_obj in \
          stream_obj_set]
    )

    link_id = 0
    for declare_set in z3_declare_set:
        offset_set = declare_set['offset_set']
        prio_set = declare_set['prio_set']

        for offset, prio in zip(offset_set, prio_set):
            name = offset['name']
            stream_id = name.split('_')[1].split('^')[0]
            stream_id = int(stream_id)
            period = stream_obj_set[stream_id].period
            instance_num = int(hyper_period/period)
            macrotick = link_obj_set[link_id].macrotick
            period_in_macrotick = int(period/macrotick)
            for i in range(0, instance_num):
                offset_of_each_instance = offset['value'] + i*period_in_macrotick
                dic = {'name': offset['name'],
                       'value': offset_of_each_instance}
                temp_z3_declare_set[link_id]['offset_set'].append(dic)
                temp_z3_declare_set[link_id]['prio_set'].append(prio)
        link_id += 1
    return temp_z3_declare_set


def write_declare_set_to_txt(result_set,
                             link_obj_set,
                             stream_obj_set,
                             solution_txt
                             ):
    time_used_in_second = result_set['time_used_in_second']
    sat_or_not = result_set['sat_or_not']
    z3_declare_set = result_set['declare_set']
    unknown_reason = result_set['unknown_reason']

    width = 52

    f = open(solution_txt, 'w')

    # 如果有可行解
    if sat_or_not == 'sat':

        z3_declare_set = _classify_declare_set(z3_declare_set, len(link_obj_set))

        # 扩展该set
        z3_declare_set = _expand_declare_set(z3_declare_set, link_obj_set, stream_obj_set)

        sorted_declare_set = _sort_declare_set(z3_declare_set)

        link_id = 0
        for declare in sorted_declare_set:
            src_node_id = link_obj_set[link_id].src_node
            dst_node_id = link_obj_set[link_id].dst_node

            divider_between_link_str = '=' * width
            f.write("%s\n" % divider_between_link_str)

            src_to_dst = '(%d, %d)' % (src_node_id, dst_node_id)
            f.write('offset and prio set at link %s, link_id: %d\n' % (src_to_dst, link_id))

            '''打印这条链路上所有的offset和prio'''
            _format_offset_set_and_prio_set(f, declare, width)

            f.write('\n')

            link_id += 1
    elif sat_or_not == 'unsat':
        f.write('unsat\n')
    elif sat_or_not == 'unknown':
        f.write('unknown\n')
        # 写入原因
        f.write('the reason for unknown result: %s\n' % unknown_reason)

    # 将求解时间写入文件
    _write_time_used_to_txt(f, time_used_in_second)
    print('the result is written into %s' % os.path.abspath(solution_txt))

    f.close()
